import sqlglot

# 常见的MySQL SQL语句示例
mysql_sql_examples = [
    # SELECT语句
    "SELECT * FROM users",
    "SELECT id, name, email FROM users WHERE id = 1",
    "SELECT COUNT(*) FROM users",
    "SELECT DISTINCT department FROM employees",
    "SELECT u.name, e.department FROM users u JOIN employees e ON u.id = e.user_id",
    "SELECT * FROM users ORDER BY created_at DESC LIMIT 10",
    "SELECT * FROM users GROUP BY department",
    "SELECT department, COUNT(*) as count FROM users GROUP BY department HAVING COUNT(*) > 5",
    # INSERT语句
    "INSERT INTO users (name, email, age) VALUES ('John', 'john@example.com', 25)",
    "INSERT INTO users VALUES (1, 'Alice', 'alice@example.com', 30)",
    # UPDATE语句
    "UPDATE users SET name = 'Bob' WHERE id = 1",
    "UPDATE users SET name = 'Alice', age = 26 WHERE email = 'alice@example.com'",
    # DELETE语句
    "DELETE FROM users WHERE id = 1",
    "DELETE FROM users WHERE age < 18",
    # CREATE DATABASE语句
    "CREATE DATABASE test",
    # CREATE TABLE语句
    "CREATE TABLE users (id INT PRIMARY KEY, name VARCHAR(50), email VARCHAR(100))",
    "CREATE TABLE employees (id INT AUTO_INCREMENT PRIMARY KEY, name VARCHAR(50), department VARCHAR(50), salary DECIMAL(10,2))",
    # ALTER TABLE语句
    "ALTER TABLE users ADD COLUMN age INT",
    "ALTER TABLE users MODIFY COLUMN name VARCHAR(100)",
    "ALTER TABLE users DROP COLUMN age",
    # DROP TABLE语句
    "DROP TABLE users",
    "DROP TABLE IF EXISTS employees",
    # CREATE INDEX语句
    "CREATE INDEX idx_users_email ON users(email)",
    "CREATE UNIQUE INDEX idx_users_id ON users(id)",
    # 子查询
    "SELECT * FROM users WHERE id IN (SELECT user_id FROM employees WHERE department = 'IT')",
    "SELECT (SELECT COUNT(*) FROM users) as total_users",
    # UNION操作
    "SELECT name FROM users UNION SELECT name FROM employees",
    "SELECT name FROM users UNION ALL SELECT name FROM employees",
    # 复杂JOIN
    "SELECT u.name, e.department, p.project_name FROM users u LEFT JOIN employees e ON u.id = e.user_id LEFT JOIN projects p ON e.id = p.employee_id",
    # CASE语句
    "SELECT name, CASE WHEN age < 18 THEN 'Minor' WHEN age >= 18 AND age < 65 THEN 'Adult' ELSE 'Senior' END as age_group FROM users",
    # 窗口函数
    "SELECT name, salary, ROW_NUMBER() OVER (PARTITION BY department ORDER BY salary DESC) as rank FROM employees",
    # WITH子句(CTE)
    "WITH active_users AS (SELECT * FROM users WHERE last_login > '2023-01-01') SELECT * FROM active_users WHERE department = 'IT'",
]
if __name__ == "__main__":
    # 解析并打印AST
    print("MySQL SQL语句及其AST解析结果:")
    print("=" * 80)

    for i, sql in enumerate(mysql_sql_examples, 1):
        print(f"\n{i}. SQL语句: {sql}")
        print("-" * 60)
        try:
            ast = sqlglot.parse_one(sql, dialect="mysql")
            print(f"AST: {repr(ast)}")
        except Exception as e:
            print(f"解析错误: {e}")
        print("=" * 80)
